# -*- coding: utf-8 -*-
"""
1. img -> JPEGImages
   xml -> Annotations

2. 将子目录中的图像和标注文件汇集到一个标准VOC标准目录中

原来的目录组织形式如下：

plane
    |
    |---J20
           |--J20_1 (xml, img)
           |--J20_2 (xml, img)
    |
    |
    |---J10
           |--J10_1 (xml, img)
           |--J10_2 (xml，img)

整理后：

plane
    |
    |--Annotations
    |--JPEGImages

重新组织以后保持原来xml文件中的类别名称不变
"""

import os
import shutil
import sys
import argparse
try:
    import xml.etree.cElementTree as ET
except ImportError:
    import xml.etree.ElementTree as ET


def reorganise_directory(root_dir, save_dir, rename=False):
    img_count = 0
    xml_count = 0
    for root, dirs, files in os.walk(root_dir):
        if files:
            old_dir_name = root
            new_dir_name = None
            if rename:
                if os.path.split(old_dir_name)[-1] == "xml":
                    new_dir_name = old_dir_name.replace(os.path.split(old_dir_name)[-1], "Annotations")
                elif os.path.split(old_dir_name)[-1] == "img":
                    new_dir_name = old_dir_name.replace(os.path.split(old_dir_name)[-1], "JPEGImages")
                if new_dir_name is not None:
                    os.rename(old_dir_name, new_dir_name)
            for file in files:
                old_file_full_path = os.path.join(root, file)
                file_path = root
                if len(file_path.split("/")) >= 3:
                    if file_path.split("/")[-3] is not None and file_path.split("/")[-2] is not None:
                        save_file_name = file_path.split("/")[-3] + "_" + file_path.split("/")[-2]+"_"+file
                    elif file_path.split("/")[-2] is not None:
                        save_file_name = file_path.split("/")[-2] + "_" + file
                    else:
                        save_file_name = file
                # if len(save_file_name.split(" ")) >= 2:
                #     save_file_name = ''.join(save_file_name.split())
                if os.path.split(old_dir_name)[-1] == "img" or os.path.split(old_dir_name)[-1] == "JPEGImages":
                    # new_file_full_path = os.path.join(save_dir, "JPEGImages", ''.join(save_file_name.split()))
                    new_file_full_path = os.path.join(save_dir, "JPEGImages", save_file_name)
                    img_count += 1
                elif os.path.split(old_dir_name)[-1] == "xml" or os.path.split(old_dir_name)[-1] == "Annotations":
                    # xml_file_full_path = old_file_full_path
                    # tree = ET.parse(xml_file_full_path)
                    # xml_root = tree.getroot()
                    # img_name_in_xml = xml_root.find('filename').text
                    # # if img_name_in_xml.endswith("jpg") or img_name_in_xml.endswith("png") or img_name_in_xml.endswith("PNG")\
                    # #         or img_name_in_xml.endswith("jpeg") or img_name_in_xml.endswith("JPG"):
                    # _, img_ext = os.path.splitext(img_name_in_xml)
                    # if not img_ext:
                    #     img_ext = '.jpg'
                    # base_img_name, _ = os.path.splitext(save_file_name)
                    #
                    # xml_root.find('filename').text = base_img_name + img_ext
                    # tree.write(xml_file_full_path)

                    new_file_full_path = os.path.join(save_dir, "Annotations", save_file_name)
                    xml_count += 1
                if new_file_full_path is not None:
                    if not os.path.exists(os.path.split(new_file_full_path)[0]):
                        os.makedirs(os.path.split(new_file_full_path)[0])

                if os.path.exists(new_file_full_path):
                    continue
                    # file_name, extend = os.path.splitext(new_file_full_path)
                    # base_name, name_id = os.path.split(file_name).split("_")
                    # new_file_full_path = base_name + "_" + str(int(name_id) + 1) + extend
                shutil.copyfile(old_file_full_path, new_file_full_path)
                if img_count % 1000 == 0:
                    print("copied images: {}".format(img_count))
                if xml_count % 1000 == 0:
                    print("copied xml: {}".format(xml_count))
    print("*"*80)
    print("Total img copied: {}".format(img_count))
    print("Total xml copied: {}".format(xml_count))


if __name__ == '__main__':

    PARSER = argparse.ArgumentParser(description="arguments")
    PARSER.add_argument("--root_dir", type=str,
                        default="/zongshiban_data/2/ccj/dataset/plane_ship_tank",
                        help="VOC dataset root directory")
    PARSER.add_argument("--save_dir", type=str,
                        default="/zongshiban_data/2/ccj/dataset/plane_ship_tank_VOC",
                        help="VOC dataset root directory")
    ARGS = PARSER.parse_args()
    ROOT_DIR = ARGS.root_dir
    SAVE_DIR = ARGS.save_dir

    reorganise_directory(ROOT_DIR, SAVE_DIR)